lang/C/sorting.and.searching.cormen.algo/Introduction to Algorithms, Third Edition/code/rbt.c

   1 // red-black tree
   2
   3 #include <stdio.h>
   4 #include <stdlib.h>
   5 #include <string.h>
   6 #include <stdarg.h>
   7
   8 //////////////////////
   9 // supplied by user //
  10 //////////////////////
  11
  12 typedef int KeyType;            // type of key
  13
  14 typedef struct {                // value related to key
  15     int stuff;
  16 } ValType;
  17
  18 // how to compare keys
  19 #define compLT(a,b) (a < b)
  20 #define compEQ(a,b) (a == b)
  21
  22 /////////////////////////////////////
  23 // implementation independent code //
  24 /////////////////////////////////////
  25
  26 typedef enum {
  27     RBT_STATUS_OK,
  28     RBT_STATUS_MEM_EXHAUSTED,
  29     RBT_STATUS_DUPLICATE_KEY,
  30     RBT_STATUS_KEY_NOT_FOUND
  31 } RbtStatus;
  32
  33 typedef enum { BLACK, RED } nodeColor;
  34
  35 typedef struct NodeTag {
  36     struct NodeTag *left;       // left child
  37     struct NodeTag *right;      // right child
  38     struct NodeTag *parent;     // parent
  39     nodeColor color;            // node color (BLACK, RED)
  40     KeyType key;                // key used for searching
  41     ValType val;                // data related to key
  42 } NodeType;
  43
  44 #define SENTINEL &sentinel      // all leafs are sentinels
  45 static NodeType sentinel = { SENTINEL, SENTINEL, 0, BLACK, 0};
  46
  47 static NodeType *root = SENTINEL; // root of red-black tree
  48
  49 static void rotateLeft(NodeType *x) {
  50
  51     // rotate node x to left
  52
  53     NodeType *y = x->right;
  54
  55     // establish x->right link
  56     x->right = y->left;
  57     if (y->left != SENTINEL) y->left->parent = x;
  58
  59     // establish y->parent link
  60     if (y != SENTINEL) y->parent = x->parent;
  61     if (x->parent) {
  62         if (x == x->parent->left)
  63             x->parent->left = y;
  64         else
  65             x->parent->right = y;
  66     } else {
  67         root = y;
  68     }
  69
  70     // link x and y
  71     y->left = x;
  72     if (x != SENTINEL) x->parent = y;
  73 }
  74
  75 static void rotateRight(NodeType *x) {
  76
  77     // rotate node x to right
  78
  79     NodeType *y = x->left;
  80
  81     // establish x->left link
  82     x->left = y->right;
  83     if (y->right != SENTINEL) y->right->parent = x;
  84
  85     // establish y->parent link
  86     if (y != SENTINEL) y->parent = x->parent;
  87     if (x->parent) {
  88         if (x == x->parent->right)
  89             x->parent->right = y;
  90         else
  91             x->parent->left = y;
  92     } else {
  93         root = y;
  94     }
  95
  96     // link x and y
  97     y->right = x;
  98     if (x != SENTINEL) x->parent = y;
  99 }
 100
 101 static void insertFixup(NodeType *x) {
 102
 103     // maintain red-black tree balance
 104     // after inserting node x
 105
 106     // check red-black properties
 107     while (x != root && x->parent->color == RED) {
 108         // we have a violation
 109         if (x->parent == x->parent->parent->left) {
 110             NodeType *y = x->parent->parent->right;
 111             if (y->color == RED) {
 112
 113                 // uncle is RED
 114                 x->parent->color = BLACK;
 115                 y->color = BLACK;
 116                 x->parent->parent->color = RED;
 117                 x = x->parent->parent;
 118             } else {
 119
 120                 // uncle is BLACK
 121                 if (x == x->parent->right) {
 122                     // make x a left child
 123                     x = x->parent;
 124                     rotateLeft(x);
 125                 }
 126
 127                 // recolor and rotate
 128                 x->parent->color = BLACK;
 129                 x->parent->parent->color = RED;
 130                 rotateRight(x->parent->parent);
 131             }
 132         } else {
 133
 134             // mirror image of above code
 135             NodeType *y = x->parent->parent->left;
 136             if (y->color == RED) {
 137
 138                 // uncle is RED
 139                 x->parent->color = BLACK;
 140                 y->color = BLACK;
 141                 x->parent->parent->color = RED;
 142                 x = x->parent->parent;
 143             } else {
 144
 145                 // uncle is BLACK
 146                 if (x == x->parent->left) {
 147                     x = x->parent;
 148                     rotateRight(x);
 149                 }
 150                 x->parent->color = BLACK;
 151                 x->parent->parent->color = RED;
 152                 rotateLeft(x->parent->parent);
 153             }
 154         }
 155     }
 156     root->color = BLACK;
 157 }
 158
 159 // insert new node (no duplicates allowed)
 160 RbtStatus rbtInsert(KeyType key, ValType val) {
 161     NodeType *current, *parent, *x;
 162
 163     // allocate node for data and insert in tree
 164
 165     // find future parent
 166     current = root;
 167     parent = 0;
 168     while (current != SENTINEL) {
 169         if (compEQ(key, current->key))
 170             return RBT_STATUS_DUPLICATE_KEY;
 171         parent = current;
 172         current = compLT(key, current->key) ?
 173             current->left : current->right;
 174     }
 175
 176     // setup new node
 177     if ((x = malloc (sizeof(*x))) == 0)
 178         return RBT_STATUS_MEM_EXHAUSTED;
 179     x->parent = parent;
 180     x->left = SENTINEL;
 181     x->right = SENTINEL;
 182     x->color = RED;
 183     x->key = key;
 184     x->val = val;
 185
 186     // insert node in tree
 187     if(parent) {
 188         if(compLT(key, parent->key))
 189             parent->left = x;
 190         else
 191             parent->right = x;
 192     } else {
 193         root = x;
 194     }
 195
 196     insertFixup(x);
 197
 198     return RBT_STATUS_OK;
 199 }
 200
 201 static void deleteFixup(NodeType *x) {
 202
 203     // maintain red-black tree balance
 204     // after deleting node x
 205
 206     while (x != root && x->color == BLACK) {
 207         if (x == x->parent->left) {
 208             NodeType *w = x->parent->right;
 209             if (w->color == RED) {
 210                 w->color = BLACK;
 211                 x->parent->color = RED;
 212                 rotateLeft (x->parent);
 213                 w = x->parent->right;
 214             }
 215             if (w->left->color == BLACK && w->right->color == BLACK) {
 216                 w->color = RED;
 217                 x = x->parent;
 218             } else {
 219                 if (w->right->color == BLACK) {
 220                     w->left->color = BLACK;
 221                     w->color = RED;
 222                     rotateRight (w);
 223                     w = x->parent->right;
 224                 }
 225                 w->color = x->parent->color;
 226                 x->parent->color = BLACK;
 227                 w->right->color = BLACK;
 228                 rotateLeft (x->parent);
 229                 x = root;
 230             }
 231         } else {
 232             NodeType *w = x->parent->left;
 233             if (w->color == RED) {
 234                 w->color = BLACK;
 235                 x->parent->color = RED;
 236                 rotateRight (x->parent);
 237                 w = x->parent->left;
 238             }
 239             if (w->right->color == BLACK && w->left->color == BLACK) {
 240                 w->color = RED;
 241                 x = x->parent;
 242             } else {
 243                 if (w->left->color == BLACK) {
 244                     w->right->color = BLACK;
 245                     w->color = RED;
 246                     rotateLeft (w);
 247                     w = x->parent->left;
 248                 }
 249                 w->color = x->parent->color;
 250                 x->parent->color = BLACK;
 251                 w->left->color = BLACK;
 252                 rotateRight (x->parent);
 253                 x = root;
 254             }
 255         }
 256     }
 257     x->color = BLACK;
 258 }
 259
 260 // delete node
 261 RbtStatus rbtErase(NodeType * z) {
 262     NodeType *x, *y;
 263
 264     if (z->left == SENTINEL || z->right == SENTINEL) {
 265         // y has a SENTINEL node as a child
 266         y = z;
 267     } else {
 268         // find tree successor with a SENTINEL node as a child
 269         y = z->right;
 270         while (y->left != SENTINEL) y = y->left;
 271     }
 272
 273     // x is y's only child
 274     if (y->left != SENTINEL)
 275         x = y->left;
 276     else
 277         x = y->right;
 278
 279     // remove y from the parent chain
 280     x->parent = y->parent;
 281     if (y->parent)
 282         if (y == y->parent->left)
 283             y->parent->left = x;
 284         else
 285             y->parent->right = x;
 286     else
 287         root = x;
 288
 289     if (y != z) {
 290         z->key = y->key;
 291         z->val = y->val;
 292     }
 293
 294
 295     if (y->color == BLACK)
 296         deleteFixup (x);
 297
 298     free (y);
 299
 300     return RBT_STATUS_OK;
 301 }
 302
 303 // find key
 304 NodeType *rbtFind(KeyType key) {
 305     NodeType *current;
 306     current = root;
 307     while(current != SENTINEL) {
 308         if(compEQ(key, current->key)) {
 309             return current;
 310         } else {
 311             current = compLT (key, current->key) ?
 312                 current->left : current->right;
 313         }
 314     }
 315     return NULL;
 316 }
 317
 318 // in-order walk of tree
 319 void rbtInorder(NodeType *p, void (callback)(NodeType *)) {
 320     if (p == SENTINEL) return;
 321     rbtInorder(p->left, callback);
 322     callback(p);
 323     rbtInorder(p->right, callback);
 324 }
 325
 326 // delete nodes depth-first
 327 void rbtDelete(NodeType *p) {
 328     if (p == SENTINEL) return;
 329     rbtDelete(p->left);
 330     rbtDelete(p->right);
 331     free(p);
 332 }
 333
 334 void displayNode(NodeType *p) {
 335     printf("%d %d\n", p->key, p->val.stuff);
 336 }
 337
 338 int main(int argc, char **argv) {
 339     int maxnum, ct;
 340     KeyType key;
 341     RbtStatus status;
 342
 343     // command-line:
 344     //
 345     //   rbt 2000
 346     //      process 2000 records
 347
 348     NodeType *p;
 349
 350     maxnum = atoi(argv[1]);
 351
 352     printf("maxnum = %d\n", maxnum);
 353     for (ct = maxnum; ct; ct--) {
 354         key = rand() % 90 + 1;
 355         if ((p = rbtFind(key)) != NULL) {
 356             if (p->val.stuff != 10*key) printf("fail val\n");
 357             status = rbtErase(p);
 358             if (status) printf("fail: status = %d\n", status);
 359         } else {
 360             ValType val;
 361             val.stuff = 10*key;
 362             status = rbtInsert(key, val);
 363             if (status) printf("fail: status = %d\n", status);
 364         }
 365     }
 366
 367     // output nodes in order
 368     rbtInorder(root, displayNode);
 369
 370     rbtDelete(root);
 371
 372     return 0;
 373 }